#!/usr/bin/env python3
#
# Silver Data mainly means TTM data
#

from numpy.core.numeric import MAXDIMS
from qiutil import *
import pandas as pd
import numpy as np
from dataclasses import dataclass

@dataclass
class TtmRoe:
    equity: pd.DataFrame
    income: pd.DataFrame
    roe: pd.DataFrame

QITTM_TTM_EQUITY_KEY = 'ttm__equity'
QITTM_TTM_INCOME_KEY ='ttm__income'
QITTM_TTM_ROE_KEY = 'ttm__roe'

def calc_equity_ttm(raw):
    # ttm = self.qd.equity_data[['STOCK_CODE', 'E_RPT_PERIOD', 'TOTAL_EQUITY_PARENT', 'ANNOUCE_DATE', 'BSPAN', 'ANN_SPAN']].copy()
    ttm = raw.equity_data[['STOCK_CODE', 'E_RPT_PERIOD', 'TOTAL_EQUITY_PARENT', 'ANNOUCE_DATE', 'E_RPT_PERIOD2', 'ANN_DATE2']].copy()
    ttm.sort_values(by=['STOCK_CODE', 'E_RPT_PERIOD']).reset_index(drop=True, inplace=True)
    ttm[['STOCK1', 'PREV_PERIOD1', 'PREV1']] = ttm[['STOCK_CODE', 'E_RPT_PERIOD', 'TOTAL_EQUITY_PARENT']].shift(1) # [['STOCK_CODE', 'E_RPT_PERIOD', 'TOTAL_EQUITY_PARENT']]
    ttm[['STOCK2', 'PREV_PERIOD2', 'PREV2']] = ttm[['STOCK_CODE', 'E_RPT_PERIOD', 'TOTAL_EQUITY_PARENT']].shift(2)
    ttm[['STOCK3', 'PREV_PERIOD3', 'PREV3']] = ttm[['STOCK_CODE', 'E_RPT_PERIOD', 'TOTAL_EQUITY_PARENT']].shift(3)
    ttm = ttm.query('STOCK_CODE == STOCK3')

    print("Before filter: " + str(ttm.shape))
    ttm = ttm[(ttm['PREV_PERIOD1'] == ttm['E_RPT_PERIOD'].map(PrevReportPeriod)) &
              (ttm['PREV_PERIOD2'] == ttm['PREV_PERIOD1'].map(PrevReportPeriod)) &
              (ttm['PREV_PERIOD3'] == ttm['PREV_PERIOD2'].map(PrevReportPeriod))]
    print("After filter: " + str(ttm.shape))
    ttm['EQUITY_TTM'] = ttm[['TOTAL_EQUITY_PARENT', 'PREV1', 'PREV2', 'PREV3']].mean(axis=1)
    return ttm[['STOCK_CODE', 'E_RPT_PERIOD',
                'TOTAL_EQUITY_PARENT', 'EQUITY_TTM','PREV1', 'PREV2', 'PREV3', 'ANNOUCE_DATE', 'E_RPT_PERIOD2', 'ANN_DATE2']]


def calc_income_ttm(raw):
    # 计算每年的总Profit
    # 计算每年的总Profit
    bool_index = raw.income_data['I_RPT_PERIOD'].str.endswith("1231") | (raw.income_data['I_RPT_PERIOD'] == "20210930")
    annual_income = raw.income_data[bool_index][['STOCK_CODE', 'I_RPT_PERIOD', 'NET_PROFIT_GM']]
    annual_income = annual_income.rename(columns={"NET_PROFIT_GM" : "YEAR_TOT_PROFIT"})
    annual_income['YEAR'] = annual_income['I_RPT_PERIOD'].str.slice(0, 4)
    annual_income = annual_income.drop(columns=['I_RPT_PERIOD'])
    # print(annual_income)

    # 将 每年的总Profit列 并入income_ttm表
    # use inner instead of left, ignore stocks without report_period ends with 1231 or 20210930
    income_ttm = pd.merge(raw.income_data, annual_income, how='inner',
                      on=["STOCK_CODE", "YEAR"], suffixes=("", "_Y"))
    # new_df2 = new_df.query('i == "left_only"')

    # 计算每个季度末对应的残余Profit 且将其存入prevq_fut_profit表
    income_ttm["PREVQ_FUT_PROFIT"] = income_ttm["YEAR_TOT_PROFIT"] - income_ttm["NET_PROFIT_GM"]
    prevq_fut_profit = income_ttm[["STOCK_CODE", "I_RPT_PERIOD", "PREVQ_FUT_PROFIT"]]
    income_ttm = income_ttm.drop(columns=["PREVQ_FUT_PROFIT"])


    # 计算每个季度末对应的前一个环比季度末， 且将对应列并入income_ttm表
    income_ttm['PREVQ_END'] = income_ttm['I_RPT_PERIOD'].map(lambda x: str(int(x)-10000))

    # 通过join income_ttm表里的PREVQ_END和prevq_fut_profit的I_RPT_PERIOD
    # 将每个季度末对应的上年利润作为“PREVQ_FUT_PROFIT”计算出来并放入income_ttm表
    income_ttm = pd.merge(income_ttm, prevq_fut_profit, how="inner",
                       left_on=["STOCK_CODE", "PREVQ_END"],
                       right_on=["STOCK_CODE", "I_RPT_PERIOD"],
                       suffixes=("", "_Z"))

    # 整理income_ttm表
    income_ttm = income_ttm.drop(columns=["I_RPT_PERIOD_Z"])

    # 计算PROFIT_TTM
    income_ttm['PROFIT_TTM'] = income_ttm['NET_PROFIT_GM'] + income_ttm['PREVQ_FUT_PROFIT']
    # print(income_ttm)

    return income_ttm[['STOCK_CODE', 'I_RPT_PERIOD', 'NET_PROFIT_GM', 'PROFIT_TTM', 'ANN_DT']]
    # print(income_ttm)
    # 暂时用所有的PROFIT_TTM作直方图，以后再细化按季度作直方图
    # sns.histplot(np.log(income_ttm["PROFIT_TTM"]), bins=30, kde=True)

def add_rpt_end_to_income(income):
    df = income.copy()
    df.sort_values(by=['STOCK_CODE', 'I_RPT_PERIOD'], inplace=True)
    df.reset_index(inplace=True, drop=True)

    df[['STOCK_CODE2', 'I_RPT_PERIOD2', 'ANN_DT2']]  = df.shift(-1)[['STOCK_CODE', 'I_RPT_PERIOD', 'ANN_DT']]
    bool_ind = df['STOCK_CODE'] != df['STOCK_CODE2']
    df.loc[bool_ind, 'I_RPT_PERIOD2'] = MAX_END_DATE
    df.drop(columns=['STOCK_CODE2'], inplace=True)
    return df

def calc_roe_ttm(income, equity):
    # 计算 roe_ttm
    print("calc_roe_ttm: ", end="")
    print(equity.columns)

    roe_ttm = pd.merge(income, equity,
                       left_on=['STOCK_CODE', 'I_RPT_PERIOD'],
                       right_on=['STOCK_CODE', 'E_RPT_PERIOD'],
                       how='inner')

    roe_ttm['ROE_TTM'] = roe_ttm['PROFIT_TTM'] / roe_ttm['EQUITY_TTM']

    roe_ttm = roe_ttm.rename(columns={'ANN_DT':'I_ANNDT', 'ANNOUCE_DATE':'E_ANNDT', 'ANN_DT2':'I_ANNDT2', 'ANN_DATE2':'E_ANNDT2'})
    conditions = [roe_ttm['I_ANNDT'] < roe_ttm['E_ANNDT'], roe_ttm['I_ANNDT'] >= roe_ttm['E_ANNDT']]
    choices = [roe_ttm['I_ANNDT'], roe_ttm['E_ANNDT']]
    choices2 = [roe_ttm['I_ANNDT2'], roe_ttm['E_ANNDT2']]
    roe_ttm['SEL_ANNDT'] = np.select(conditions, choices)
    roe_ttm['SEL_ANNDT2'] = np.select(conditions, choices2)
    # roe_ttm = roe_ttm.drop(columns=['INCOME_ANN_DT', 'EQUITY_ANN_DT'])
    return roe_ttm

def load_ttm_roe_from_h5file(data_file=DEFAULT_SILVER_FILE):
    store = pd.HDFStore(data_file, mode='r+')
    equity = store.get(QITTM_TTM_EQUITY_KEY)
    income = store.get(QITTM_TTM_INCOME_KEY)
    roe = store.get(QITTM_TTM_ROE_KEY)
    ttmroe = TtmRoe(equity, income, roe)
    return ttmroe

def load_ttm_roe_from_h5store(store):
    equity = store.get(QITTM_TTM_EQUITY_KEY)
    income = store.get(QITTM_TTM_INCOME_KEY)
    roe = store.get(QITTM_TTM_ROE_KEY)
    ttmroe = TtmRoe(equity, income, roe)
    return ttmroe


def save_ttm_roe(ttmroe, store):
    ttmroe.roe.to_hdf(store, QITTM_TTM_ROE_KEY, mode='a')
    ttmroe.equity.to_hdf(store, QITTM_TTM_EQUITY_KEY, mode='a')
    ttmroe.income.to_hdf(store, QITTM_TTM_INCOME_KEY, mode='a')

def cleanse_roe(ttm):
    ttm.roe = ttm.roe[ttm.roe['ROE_TTM'] != -np.inf]
    ttm.roe.ROE_TTM.mean()

    for ratio in CLEANSE_RATIO:
        deviation = np.abs(ttm.roe.ROE_TTM - ttm.roe.ROE_TTM.mean())
        roe_filter = deviation <= (ratio * ttm.roe.ROE_TTM.std())
        ttm.roe = ttm.roe[roe_filter]


def create_ttm_roe(raw_data,
                   equity_ttm_func=calc_equity_ttm,
                   income_ttm_func=calc_income_ttm) -> TtmRoe:
    print("create_ttm_roe: ", end="")
    print(raw_data.equity_data.columns)
    equity = equity_ttm_func(raw_data)

    print(equity.columns)
    income = income_ttm_func(raw_data)
    income = add_rpt_end_to_income(income)
    roe = calc_roe_ttm(income, equity)
    ttm = TtmRoe(equity, income, roe)
    cleanse_roe(ttm)
    return ttm
